import MySQLdb
import pyarrow as pa


MYSQL_TO_ARROW_TYPES = {
    MySQLdb.FIELD_TYPE.TINY: pa.int8(),  # TINYINT
    MySQLdb.FIELD_TYPE.SHORT: pa.int16(),  # SMALLINT
    MySQLdb.FIELD_TYPE.LONG: pa.int32(),  # INTEGER
    MySQLdb.FIELD_TYPE.FLOAT: pa.float32(),  # FLOAT
    MySQLdb.FIELD_TYPE.DOUBLE: pa.float64(),  # DOUBLE
    MySQLdb.FIELD_TYPE.LONGLONG: pa.int64(),  # BIGINT
    MySQLdb.FIELD_TYPE.INT24: pa.int32(),  # MEDIUMINT
    MySQLdb.FIELD_TYPE.YEAR: pa.int16(),  # YEAR
    MySQLdb.FIELD_TYPE.DECIMAL: pa.decimal128(precision=38, scale=4),  # DECIMAL
    MySQLdb.FIELD_TYPE.NEWDECIMAL: pa.decimal128(precision=38, scale=4),  # NEWDECIMAL
    MySQLdb.FIELD_TYPE.TIMESTAMP: pa.timestamp("us"),  # TIMESTAMP
    MySQLdb.FIELD_TYPE.DATE: pa.date32(),  # DATE
    MySQLdb.FIELD_TYPE.TIME: pa.time32("s"),  # TIME
    MySQLdb.FIELD_TYPE.DATETIME: pa.timestamp("us"),  # DATETIME
    MySQLdb.FIELD_TYPE.STRING: pa.string(),  # CHAR, VARCHAR, TEXT
    MySQLdb.FIELD_TYPE.VAR_STRING: pa.string(),  # VARCHAR
    MySQLdb.FIELD_TYPE.BLOB: pa.binary(),  # BLOB
}


PSYCOPG_TO_ARROW_TYPES = {
    16: pa.string(),  # CHAR, VARCHAR, TEXT
    17: pa.string(),  # BYTEA
    20: pa.int64(),  # BIGINT
    21: pa.int8(),  # SMALLINT
    23: pa.int32(),  # INTEGER
    25: pa.string(),  # TEXT
    700: pa.float32(),  # REAL
    701: pa.float64(),  # DOUBLE PRECISION
    1042: pa.string(),  # CHAR
    1043: pa.string(),  # VARCHAR
    1082: pa.date32(),  # DATE
    1114: pa.timestamp("us"),  # TIMESTAMP WITHOUT TIME ZONE
    1184: pa.timestamp("us"),  # TIMESTAMP WITH TIME ZONE
    1700: pa.decimal128(precision=38, scale=4),  # NUMERIC
}
